import hashlib
import os.path
import re
from datetime import datetime

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from hetong.db.models.knowledge_base_model import KnowledgeBean
from hetong.doc_hand.utils import add_file_to_db
from hetong.qwen.index import extract_info
from hetong.qwen.online import extract_info as online_info
from hetong.util.ziputil import ZipFileHand
from hetong.util.ftp_util import download_file, upload_file

from hetong.config import *
from configs import logger
from hetong.db.knowledge_base_repository import list_kbs_from_db
import json
import pandas as pd
from hetong.db.lib_job.job_repository import update_task, update_percentage


def read_file_md5(file_path):
    with open(file_path, 'rb') as f:
        return hashlib.md5(f.read()).hexdigest()


class DocHand:
    def __init__(self, path=None, task_id=None):
        try:
            update_task(task_id, status=2, start_time=True, percentage=10)
            self.start_time = datetime.now()
            self.path = path
            if path:
                self.file_name = os.path.basename(path)
            self.task_id = task_id
            self.csv_data = []
        except Exception as e:
            logger.exception(e)
            update_task(task_id, status=3, end_time=True, reasons=str(e))

    def init_data(self):
        try:
            self.init_path()
            self.load_doc()
        except Exception as e:
            logger.exception(e)
            update_task(self.task_id, status=3, end_time=True, reasons=str(e))

    def init_path(self):
        try:
            '''
            解析文件夹
            '''
            local_path = self.path
            if use_ftp:
                local_path = self.check_path(local_path)
            if os.path.exists(local_path):
                self.file_md5 = read_file_md5(local_path)
            if local_path.endswith('zip') or local_path.endswith('ZIP'):
                base_out_path = os.path.dirname(local_path)
                out_path = os.path.join(base_out_path, self.file_md5)
                if not os.path.exists(out_path):
                    os.mkdir(out_path)
                zip_hand = ZipFileHand(local_path, out_path)
                zip_hand.hand_zip_file()
                out_dir = zip_hand.out_dir
            else:
                out_dir = local_path
            self.out_dir = out_dir
        except Exception as e:
            logger.exception(e)

    def check_path(self, path):
        try:
            file_name = os.path.basename(path)
            local_path = '%s/%s' % (config.out_dir, file_name)
            if not os.path.exists(os.path.dirname(local_path)):
                os.mkdir(os.path.dirname(local_path))
            if not os.path.exists(local_path):
                download_file(local_path, path, ftp_config=config.ftp_config)
                logger.debug("下载文件成功")
                return local_path
            return local_path
        except Exception as e:
            logger.exception(e)
            return None

    def load_hetong_doc(self, path):
        '''
        加载合同   这一层是合同的实际位置
        '''
        kb_name = os.path.basename(path)
        files = []
        if os.path.isdir(path):
            for hetong_file in os.listdir(path):
                hetong = os.path.join(path, hetong_file)
                if os.path.isfile(hetong):
                    files.append(hetong)
        # 最后一层才是合同
        if len(files) > 0:
            for item in add_file_to_db(self.task_id, knowledge_base_name=kb_name, file_paths=files,
                                       chunk_size=800):
                print(item)
                # 持续更新下进度
            update_percentage(self.task_id)

    def load_doc(self):
        '''
        加载文档并解析，用知识库的方式
        第一层文件夹为合同编号，同时当知识库的名称
        '''
        if self.out_dir:
            for item in os.listdir(os.path.join(self.out_dir)):
                # 第一层为合同文件夹
                file_path = os.path.join(self.out_dir, item)
                if os.path.isdir(file_path):
                    # 第二层位合同实际的文件夹，如果还有压缩包，目前暂时未处理
                    for child in os.listdir(file_path):
                        self.load_hetong_doc(os.path.join(file_path, child))

    def get_item(self, data, query):
        try:
            j_value = ''
            json_blocks = re.findall(r'```json(.+?)```', data, re.DOTALL)
            for block in json_blocks:
                try:
                    j_data = json.loads(block)
                    if j_data['key'] == query:
                        if len(j_data) > len(j_value):
                            j_value = j_data['value']
                except json.JSONDecodeError:
                    print(f"无法解析的JSON数据块: {block}")
            # result = dd.split('```\n')[0].replace('```json', '').replace('\n', '').replace(' ', '').replace('```', '')
            # j_data = json.loads(result)
            # if j_data['key']==query:
            #     return j_data['value']
            return j_value
        except Exception as e:
            logger.exception(e)
        return ''

    def hand_single_hetong(self, kb: KnowledgeBean):
        '''
        按单个合同生成数据值
        '''
        from hetong.doc_hand.base import KBServiceFactory
        kb_name = kb.kb_name
        hetong_data = {'合同编号': kb_name}
        for key in keys:
            try:
                query = key['name']
                other_keys = key['keys']
                search_keys = [item for item in other_keys]
                search_keys.append(query)
                kb = KBServiceFactory.get_service(kb_name=kb_name, vector_store_type=DEFAULT_VS_TYPE)
                result = []
                hetong_data[query] = ''
                for search_key in search_keys:
                    try:
                        docs = kb.search_docs(search_key)
                        for doc in docs:
                            result.append(doc)
                    except Exception as e:
                        logger.exception(e)
                data = [item[0].page_content for item in docs]
                if len(data) == 0:
                    continue
                if use_online_mode:
                    res = online_info('\n'.join(data), query)
                else:
                    res = extract_info('\n'.join(data), query)
                hetong_data[query] = self.get_item(res, query)
            except Exception as e:
                logger.exception(e)
        self.csv_data.append(hetong_data)
        update_percentage(self.task_id)

    def create_csv(self, path):

        result = self.csv_data
        df = pd.DataFrame(result)
        print('本次数据为:', self.csv_data)
        df.to_csv(path, index=False, sep=',')
        logger.debug('本次合同文件生成完成', path)

    def upload_file(self, local_path):
        # 上传结果文件到指定服务器的指定位置：
        base_name = os.path.basename(local_path).split(".csv")[0]
        try:
            remote_path = '%s/results/%s_%s' % (sftp_remote_path, base_name, os.path.basename(local_path))
            upload_file(local_path, remote_path, ftp_config=ftp_config)
            logger.debug("本次文件上传成功:{}", remote_path)
            return remote_path
        except Exception as e:
            logger.exception(e)
        return None

    def gen_result(self):
        '''
        生成结果
        '''
        try:
            # 获取本次任务结果
            kbs = list_kbs_from_db(task_id=self.task_id)
            for kb in kbs:
                self.hand_single_hetong(kb)
            tt = datetime.now().strftime('%Y%m%d%H%M%S')
            result_path = out_dir + '/hetong/%s_%s.csv' % (tt, str(self.task_id))
            if not os.path.exists(os.path.dirname((result_path))):
                os.makedirs(os.path.dirname((result_path)))
            # 生成csv文件
            self.create_csv(result_path)
            # 写回任务
            if use_ftp:
                result_path = self.upload_file(result_path)
            # 写回文件
            update_task(self.task_id, remote_path=result_path, status=1, end_time=True, percentage=100)
            logger.debug("本次耗时:{}", datetime.now() - self.start_time)
        except Exception as e:
            logger.exception(e)
            update_task(self.task_id, status=3, end_time=True, reasons=str(e))


if __name__ == '__main__':
    dd = DocHand(task_id=1257)
    dd.gen_result()
